import os
import argparse
from typing import List, Dict, Any, Optional, Tuple
from collections import defaultdict
import time
from threading import Thread, Lock
from werkzeug.http import parse_range_header
from flask import Flask, abort, jsonify, request, make_response, send_file
from flask_cors import CORS


class FileServerConfig:
    """Configuration class for file server."""

    def __init__(self) -> None:
        self.json_as_ascii = False
        self.jsonify_mimetype = "application/json; charset=utf-8"
        self.base_dir = os.path.abspath(".")
        self.port = 8000
        self.max_bytes_per_ip = 1024 * 1024 * 1024  # 1GB
        self.api_path = "/file"
        self.tracker_reset_interval = 3600  # 1 hour


class IPTracker:
    """Thread-safe IP traffic tracker with rate limiting."""

    def __init__(self, max_bytes_per_ip: int) -> None:
        self._tracker = defaultdict(int)
        self._lock = Lock()
        self.max_bytes_per_ip = max_bytes_per_ip

    def reset(self) -> None:
        """Reset all IP counters."""
        with self._lock:
            self._tracker.clear()

    def increment(self, ip: str, bytes_count: int) -> None:
        """Increment IP's byte counter."""
        with self._lock:
            self._tracker[ip] += bytes_count

    def get_remaining_quota(self, ip: str) -> int:
        """Get remaining quota for IP."""
        with self._lock:
            return max(0, self.max_bytes_per_ip - self._tracker.get(ip, 0))

    def is_over_limit(self, ip: str, additional_bytes: int = 0) -> bool:
        """Check if IP is over limit with additional bytes."""
        with self._lock:
            return (self._tracker.get(ip, 0) + additional_bytes) > self.max_bytes_per_ip


class FileServer:
    """Main file server application."""

    def __init__(self, config: FileServerConfig) -> None:
        self.config = config
        self.app = Flask(__name__)
        self._configure_flask()
        self.ip_tracker = IPTracker(config.max_bytes_per_ip)
        self._start_background_tasks()
        self._register_routes()

    def _configure_flask(self) -> None:
        """Configure Flask app settings."""
        self.app.config.update({
            "JSON_AS_ASCII": self.config.json_as_ascii,
            "JSONIFY_MIMETYPE": self.config.jsonify_mimetype,
        })
        CORS(self.app)

    def _start_background_tasks(self) -> None:
        """Start background threads."""
        def reset_tracker_periodically():
            while True:
                time.sleep(self.config.tracker_reset_interval)
                self.ip_tracker.reset()

        tracker_thread = Thread(target=reset_tracker_periodically)
        tracker_thread.daemon = True
        tracker_thread.start()

    def _register_routes(self) -> None:
        """Register all Flask routes."""
        self.app.before_request(self._track_ip_usage)
        self.app.after_request(self._add_rate_headers)

        self.app.route(f"{self.config.api_path}/", defaults={"path": ""})(
            self.app.route(f"{self.config.api_path}/<path:path>")(self._serve_file)
        )

        self.app.errorhandler(403)(self._handle_forbidden)
        self.app.errorhandler(404)(self._handle_not_found)
        self.app.errorhandler(429)(self._handle_too_many_requests)

    def _get_rate_headers(self, ip: str) -> Dict[str, str]:
        """Generate rate limit headers."""
        return {
            "X-RateLimit-Limit": str(self.config.max_bytes_per_ip),
            "X-RateLimit-Remaining": str(self.ip_tracker.get_remaining_quota(ip)),
        }

    def _track_ip_usage(self) -> None:
        """Check IP usage before request."""
        if request.endpoint != "_serve_file":
            return

        client_ip = request.remote_addr
        if self.ip_tracker.is_over_limit(client_ip):
            response = jsonify({
                "error": "Rate limit exceeded",
                "message": f"Data limit exceeded (max {self.config.max_bytes_per_ip/1024/1024}MB per hour)",
            })
            response.headers.update(self._get_rate_headers(client_ip))
            abort(429, response=response)

    def _add_rate_headers(self, response) -> Any:
        """Add rate limit headers to responses."""
        if request.endpoint == "_serve_file":
            response.headers.update(self._get_rate_headers(request.remote_addr))
        return response

    @staticmethod
    def _is_safe_path(base_dir: str, requested_path: str) -> bool:
        """Verify requested path is within base directory."""
        base_dir = os.path.realpath(base_dir)
        requested_path = os.path.realpath(requested_path)
        return requested_path.startswith(base_dir)

    def _get_file_info(self, base_url: str, rel_path: str, full_path: str) -> Dict[str, Any]:
        """Get file/directory metadata."""
        stat = os.stat(full_path)
        return {
            "name": os.path.basename(full_path),
            "url": (
                f"{base_url}{self.config.api_path}/{rel_path}"
                if rel_path
                else f"{base_url}{self.config.api_path}/{os.path.basename(full_path)}"
            ),
            "is_dir": os.path.isdir(full_path),
            "size": stat.st_size if os.path.isfile(full_path) else 0,
            "modified": stat.st_mtime,
        }

    def _get_file_list(self, base_url: str, rel_path: str, dir_path: str) -> List[Dict[str, Any]]:
        """Get directory listing with formatted URLs."""
        try:
            files = os.listdir(dir_path)
            return [
                self._get_file_info(
                    base_url,
                    os.path.join(rel_path, f) if rel_path else f,
                    os.path.join(dir_path, f),
                )
                for f in sorted(
                    files,
                    key=lambda x: (not os.path.isdir(os.path.join(dir_path, x)), x.lower())
                )
            ]
        except PermissionError as e:
            raise PermissionError(f"Access denied to directory: {dir_path}") from e

    def _handle_range_request(
        self, file_path: str, client_ip: str
    ) -> Optional[Tuple[Any, int]]:
        """Handle HTTP Range requests (partial content)."""
        range_header = request.headers.get("Range")
        if not range_header:
            return None

        full_size = os.path.getsize(file_path)
        try:
            ranges = parse_range_header(range_header, full_size)
            if not ranges or not ranges.ranges:
                abort(416)  # Requested range not satisfiable

            start, end = ranges.ranges[0]
            start = start if start is not None else 0
            end = end if end is not None else full_size - 1

            if start >= full_size or end >= full_size or start > end:
                abort(416)

            content_length = end - start + 1
            if self.ip_tracker.is_over_limit(client_ip, content_length):
                abort(429)

            self.ip_tracker.increment(client_ip, content_length)
            response = make_response(send_file(
                file_path,
                as_attachment=False,
                conditional=True,
            ))
            response.headers["Content-Range"] = f"bytes {start}-{end}/{full_size}"
            response.status_code = 206
            return response

        except ValueError as e:
            abort(416)

    def _serve_file(self, path: str = "") -> Any:
        """Main endpoint to serve files/directories."""
        requested_path = os.path.join(self.config.base_dir, path)

        if not self._is_safe_path(self.config.base_dir, requested_path):
            abort(403)

        client_ip = request.remote_addr

        if os.path.isdir(requested_path):
            index_path = os.path.join(requested_path, "index.html")
            if os.path.exists(index_path):
                return send_file(index_path)

            try:
                file_list = self._get_file_list(
                    request.url_root.rstrip("/"),
                    path,
                    requested_path,
                )
                return jsonify({
                    "path": path,
                    "files": file_list,
                    "remaining_quota": self.ip_tracker.get_remaining_quota(client_ip),
                })
            except PermissionError:
                abort(403)

        elif os.path.isfile(requested_path):
            try:
                # Handle range requests first
                range_response = self._handle_range_request(requested_path, client_ip)
                if range_response:
                    return range_response

                # Full file request
                full_size = os.path.getsize(requested_path)
                if self.ip_tracker.is_over_limit(client_ip, full_size):
                    abort(429)

                self.ip_tracker.increment(client_ip, full_size)
                return send_file(requested_path, as_attachment=False, conditional=True)

            except (PermissionError, FileNotFoundError):
                abort(403 if isinstance(e, PermissionError) else 404)

    def _handle_forbidden(self, e) -> Tuple[Any, int]:
        response = jsonify({"error": "Permission denied"})
        response.headers.update(self._get_rate_headers(request.remote_addr))
        return response, 403

    def _handle_not_found(self, e) -> Tuple[Any, int]:
        response = jsonify({"error": "File not found"})
        response.headers.update(self._get_rate_headers(request.remote_addr))
        return response, 404

    def _handle_too_many_requests(self, e) -> Tuple[Any, int]:
        response = jsonify({
            "error": "Rate limit exceeded",
            "message": str(e.description),
        })
        response.headers.update(self._get_rate_headers(request.remote_addr))
        return response, 429

    def run(self, host: str = "0.0.0.0", port: Optional[int] = None) -> None:
        """Run the Flask server."""
        self.app.run(host=host, port=port or self.config.port)


def parse_args() -> argparse.Namespace:
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description="Flask File Server with Rate Limiting")
    parser.add_argument(
        "--port",
        type=int,
        default=8000,
        help="Port to serve on (default: 8000)",
    )
    parser.add_argument(
        "--directory",
        type=str,
        default=".",
        help="Directory to serve (default: current directory)",
    )
    parser.add_argument(
        "--quota",
        type=int,
        default=1024 * 1024 * 1024,
        help="Data quota per IP in bytes (default: 35MB)",
    )
    return parser.parse_args()


def main() -> None:
    """Entry point for the file server."""
    args = parse_args()

    if not os.path.isdir(args.directory):
        raise FileNotFoundError(f"Directory '{args.directory}' does not exist")

    config = FileServerConfig()
    config.port = args.port
    config.base_dir = os.path.abspath(args.directory)
    config.max_bytes_per_ip = args.quota

    server = FileServer(config)
    server.run()


if __name__ == "__main__":
    main()